"""
ChatTTS 语音合成引擎
完全本地运行的开源TTS,音质自然,适合虚拟主播
GitHub: https://github.com/2noise/ChatTTS
"""

import os
import time
import wave
import uuid
import numpy as np
from typing import List, Optional, Tuple
from pathlib import Path

try:
    import ChatTTS
    import torch
    CHATTTS_AVAILABLE = True
except ImportError:
    CHATTTS_AVAILABLE = False

from pydub import AudioSegment

from ..core.models import LipSyncData
from ..utils.logger import get_logger
from ..utils.temp_file_manager import get_temp_file_manager


class ChatTTSEngine:
    """ChatTTS 本地语音合成引擎"""

    def __init__(
        self,
        device: str = "cpu",  # 使用设备: cpu 或 cuda
        compile: bool = False,  # 是否编译模型(加速,但启动慢)
        source: str = "huggingface",  # 模型来源: huggingface 或 local
        local_path: Optional[str] = None,  # 本地模型路径
        temperature: float = 0.3,  # 温度参数(0.1-1.0,越高越随机)
        top_p: float = 0.7,  # Top-P采样
        top_k: int = 20,  # Top-K采样
        enable_refine: bool = True,  # 是否启用文本优化
        speaker_seed: Optional[int] = None,  # 音色种子(用于固定音色)
        speaker_file: Optional[str] = None  # 预设音色文件路径(.pt格式)
    ):
        """
        初始化ChatTTS引擎
        
        Args:
            device: 运行设备(cpu/cuda)
            compile: 是否编译模型
            source: 模型来源
            local_path: 本地模型路径
            temperature: 温度参数
            top_p: Top-P采样
            top_k: Top-K采样
            enable_refine: 是否启用文本优化
            speaker_seed: 音色种子(None则每次随机,设置固定值可固定音色)
            speaker_file: 预设音色文件路径(.pt格式,优先级高于speaker_seed)
        """
        self.logger = get_logger()
        
        # 检查ChatTTS是否已安装
        if not CHATTTS_AVAILABLE:
            raise ImportError(
                "ChatTTS未安装! 请运行: pip install ChatTTS torch torchaudio"
            )
        
        self.device = device
        self.compile = compile
        self.source = source
        self.local_path = local_path
        self.temperature = temperature
        self.top_p = top_p
        self.top_k = top_k
        self.enable_refine = enable_refine
        self.speaker_seed = speaker_seed
        self.speaker_file = speaker_file
        
        # ChatTTS实例
        self.chat = None
        
        # 音色embedding(spk_emb)
        self.spk_emb = None
        
        # 临时文件管理
        self.temp_manager = get_temp_file_manager()
        self.temp_dir = self.temp_manager.get_temp_dir('audio')
        
        # 初始化模型
        self._initialize_model()
    
    def _initialize_model(self):
        """初始化ChatTTS模型"""
        try:
            self.logger.info("正在初始化ChatTTS引擎...")
            self.logger.info(f"  设备: {self.device}")
            self.logger.info(f"  模型来源: {self.source}")
            
            # 创建ChatTTS实例
            self.chat = ChatTTS.Chat()
            
            # 加载模型
            self.logger.info("正在加载ChatTTS模型(首次运行会下载模型,请耐心等待)...")
            
            # 使用新版API加载模型
            self.chat.load(
                compile=self.compile,
                source=self.source,
                custom_path=self.local_path if self.source == "local" else None
            )
            
            # 初始化音色
            self._initialize_speaker()
            
            self.logger.info("✓ ChatTTS引擎初始化完成!")
            self.logger.info("  特性: 本地运行 | 完全免费 | 音质自然")
            
        except Exception as e:
            self.logger.error(f"ChatTTS初始化失败: {e}")
            self.logger.error("请确保已安装: pip install ChatTTS torch torchaudio")
            raise
    
    def _initialize_speaker(self):
        """初始化音色设置"""
        try:
            # 优先使用预设音色文件
            if self.speaker_file and os.path.exists(self.speaker_file):
                self.logger.info(f"正在加载预设音色: {self.speaker_file}")
                # 加载音色并转移到正确的设备
                self.spk_emb = torch.load(self.speaker_file, weights_only=True, map_location=self.device)
                
                # 确保音色在正确的设备上
                if self.device == 'cuda' and not self.spk_emb.is_cuda:
                    self.spk_emb = self.spk_emb.cuda()
                elif self.device == 'cpu' and self.spk_emb.is_cuda:
                    self.spk_emb = self.spk_emb.cpu()
                    
                self.logger.info(f"✓ 预设音色加载成功 (设备: {self.spk_emb.device})")
            
            # 其次使用固定种子生成音色
            elif self.speaker_seed is not None:
                self.logger.info(f"使用固定音色种子: {self.speaker_seed}")
                torch.manual_seed(self.speaker_seed)
                self.spk_emb = self.chat.sample_random_speaker()
                self.logger.info("✓ 固定音色生成成功")
            
            # 默认随机音色
            else:
                self.logger.info("使用随机音色")
                self.spk_emb = self.chat.sample_random_speaker()
                self.logger.info("✓ 随机音色生成成功")
                
        except Exception as e:
            self.logger.warning(f"音色初始化失败,将使用默认设置: {e}")
            self.spk_emb = None
    
    def save_current_speaker(self, filepath: str):
        """保存当前音色到文件
        
        Args:
            filepath: 保存路径(.pt格式)
        """
        if self.spk_emb is None:
            self.logger.error("当前没有可保存的音色")
            return
        
        try:
            torch.save(self.spk_emb, filepath)
            self.logger.info(f"✓ 音色已保存到: {filepath}")
        except Exception as e:
            self.logger.error(f"音色保存失败: {e}")
    
    def synthesize(self, text: str, speaker: Optional[str] = None) -> Tuple[Optional[AudioSegment], List[LipSyncData]]:
        """
        合成语音
        
        Args:
            text: 要合成的文本
            speaker: 说话人(可选,用于声音克隆)
            
        Returns:
            (音频对象, 口型同步数据列表)
        """
        try:
            self.logger.debug(f"正在使用ChatTTS合成语音: {text[:50]}...")
            
            # 调试：检查音色是否正确加载
            if self.spk_emb is not None:
                # ChatTTS 0.2.4+ 兼容性：spk_emb可能是tensor或str
                if isinstance(self.spk_emb, torch.Tensor):
                    self.logger.debug(f"使用音色embedding (Tensor), shape: {self.spk_emb.shape}, device: {self.spk_emb.device}")
                else:
                    self.logger.debug(f"使用音色embedding (Type: {type(self.spk_emb).__name__})")
            else:
                self.logger.warning("警告: spk_emb为None, 将使用随机音色！")
            
            # 文本清理: ChatTTS不支持某些特殊字符
            text = self._clean_text(text)
            self.logger.debug(f"清理后的文本: {text}")
            
            # 准备推理参数 - 强制使用初始化的音色
            params_infer_code = ChatTTS.Chat.InferCodeParams(
                spk_emb=self.spk_emb,  # 使用初始化的音色（关键！）
                temperature=self.temperature,
                top_P=self.top_p,
                top_K=self.top_k,
            )
            
            # 生成音频
            self.logger.debug("正在生成音频...")
            start_time = time.time()
            
            wavs = self.chat.infer(
                [text],
                params_infer_code=params_infer_code,
                use_decoder=True  # 使用decoder模式,音质更好且更稳定
            )
            
            synthesis_time = time.time() - start_time
            self.logger.debug(f"音频生成耗时: {synthesis_time:.2f}秒")
            
            if wavs is None or len(wavs) == 0:
                self.logger.error("ChatTTS合成失败: 无音频输出")
                return None, []
            
            # 获取音频数据
            audio_data = wavs[0]
            
            # 保存为WAV文件
            temp_filename = f"chattts_{uuid.uuid4().hex[:8]}.wav"
            temp_file_path = self.temp_dir / temp_filename
            self.temp_manager.track_file(temp_file_path)
            
            # 写入WAV文件
            sample_rate = 24000  # ChatTTS默认采样率
            self._save_wav(str(temp_file_path), audio_data, sample_rate)
            
            # 加载为AudioSegment
            audio = AudioSegment.from_wav(str(temp_file_path))
            
            # 提取口型同步数据（优化：使用更高帧率提升精度）
            lipsync_data = self._extract_lipsync_data(str(temp_file_path), fps=60)  # 提升到60fps以匹配渲染帧率
            
            # 清理临时文件
            try:
                self.temp_manager.cleanup_file(temp_file_path)
            except Exception as e:
                self.logger.warning(f"清理临时文件失败: {e}")
            
            self.logger.info(f"✓ ChatTTS语音合成成功,时长: {len(audio)/1000:.2f}秒, 口型帧数: {len(lipsync_data)}")
            
            return audio, lipsync_data
            
        except Exception as e:
            self.logger.error(f"ChatTTS语音合成异常: {e}")
            import traceback
            self.logger.error(traceback.format_exc())
            return None, []
    
    def _clean_text(self, text: str) -> str:
        """
        清理文本,移除ChatTTS不支持的字符
        
        Args:
            text: 原始文本
            
        Returns:
            清理后的文本
        """
        # ChatTTS不支持的字符映射
        replacements = {
            '！': '!',   # 全角感叹号 -> 半角
            '？': '?',   # 全角问号 -> 半角
            '，': ',',   # 全角逗号 -> 半角
            '。': '.',   # 中文句号 -> 英文句号
            '；': ';',   # 全角分号 -> 半角
            '：': ':',   # 全角冒号 -> 半角
            '“': '"',  # 左双引号
            '”': '"',  # 右双引号
            '‘': "'",  # 左单引号
            '’': "'",  # 右单引号
            '《': '<',  # 左书名号
            '》': '>',  # 右书名号
            '（': '(',  # 全角左括号
            '）': ')',  # 全角右括号
            '、': ',',  # 顿号 -> 逗号
        }
        
        cleaned = text
        for old_char, new_char in replacements.items():
            cleaned = cleaned.replace(old_char, new_char)
        
        return cleaned
    
    def _refine_text(self, text: str) -> str:
        """
        文本优化(添加韵律标记)
        
        Args:
            text: 原始文本
            
        Returns:
            优化后的文本
        """
        # ChatTTS支持的韵律标记:
        # [oral_X] - 口语化程度 (0-9)
        # [laugh_X] - 笑声 (0-2)
        # [break_X] - 停顿 (0-7)
        # [speed_X] - 语速 (0-9)
        
        # 简单处理: 在标点符号处添加停顿
        refined = text
        
        # 在句号、问号、感叹号后添加较长停顿
        for punct in ['。', '?', '!', '?', '!']:
            refined = refined.replace(punct, f'{punct}[break_4]')
        
        # 在逗号、顿号后添加短停顿
        for punct in ['，', '、', ',']:
            refined = refined.replace(punct, f'{punct}[break_2]')
        
        return refined
    
    def _save_wav(self, filepath: str, audio_data: np.ndarray, sample_rate: int):
        """
        保存为WAV文件
        
        Args:
            filepath: 文件路径
            audio_data: 音频数据(numpy数组)
            sample_rate: 采样率
        """
        # 转换为int16格式
        audio_int16 = (audio_data * 32767).astype(np.int16)
        
        # 写入WAV文件
        with wave.open(filepath, 'wb') as wf:
            wf.setnchannels(1)  # 单声道
            wf.setsampwidth(2)  # 16-bit
            wf.setframerate(sample_rate)
            wf.writeframes(audio_int16.tobytes())
    
    def _extract_lipsync_data(self, audio_file: str, fps: int = 30) -> List[LipSyncData]:
        """
        从音频文件提取口型同步数据
        
        Args:
            audio_file: 音频文件路径
            fps: 帧率（默认30fps，降低数据量减少延迟）
            
        Returns:
            口型同步数据列表
        """
        try:
            # 打开WAV文件
            with wave.open(audio_file, 'rb') as wav_file:
                # 获取音频参数
                sample_rate = wav_file.getframerate()
                n_frames = wav_file.getnframes()
                audio_data = wav_file.readframes(n_frames)
                
                # 转换为numpy数组
                audio_array = np.frombuffer(audio_data, dtype=np.int16)
                
                # 计算每帧的时间间隔
                frame_interval = 1.0 / fps
                samples_per_frame = int(sample_rate * frame_interval)
                
                lipsync_data = []
                
                # 按帧分析音频
                for i in range(0, len(audio_array), samples_per_frame):
                    frame_data = audio_array[i:i + samples_per_frame]
                    
                    if len(frame_data) == 0:
                        break
                    
                    # 计算RMS音量（优化：使用平方根而非平均值）
                    rms = np.sqrt(np.mean(frame_data.astype(np.float32) ** 2))
                    max_amplitude = 32768.0  # 16-bit音频最大值
                    volume = min(rms / max_amplitude, 1.0)
                    
                    # 将音量映射到嘴巴张开度(0-1)
                    # 优化：使用非线性映射，增强低音量的可见度
                    # 使用幂函数曲线让口型更明显
                    if volume > 0.01:  # 过滤噪音
                        # 幂函数映射：让小音量也能有明显口型
                        mouth_open = min(np.power(volume * 3.0, 0.6), 1.0)
                    else:
                        mouth_open = 0.0
                    
                    # 添加口型数据
                    timestamp = i / sample_rate
                    lipsync_data.append(
                        LipSyncData(
                            timestamp=timestamp,
                            volume=volume,
                            mouth_open=mouth_open
                        )
                    )
                
                self.logger.debug(f"提取了 {len(lipsync_data)} 帧口型数据 (fps={fps})")
                return lipsync_data
                
        except Exception as e:
            self.logger.error(f"提取口型数据失败: {e}")
            return []
    
    def get_audio_duration(self, audio: Optional[AudioSegment] = None) -> float:
        """
        获取音频时长(秒)
        
        Args:
            audio: 音频对象
            
        Returns:
            音频时长(秒)
        """
        if audio is None:
            return 0.0
        
        return len(audio) / 1000.0
